module core.llm.llamaimport;

/**
 * 以 llama.cpp/example/main.cpp 为例，需要引入
 *
 * struct llama_context
 * struct llama_model
 * struct gpt_params
 * struct llama_token
 * struct llama_sampling_params
 * struct llama_context_params
 *
 * std::vector
 * std::ostringstream
 * std::string
 * std::mt19937
 * std::tie
 *
 * llama_backend_init()
 * llama_init_from_gpt_params()
 * llama_context_params_from_gpt_params()
 * llama_new_context_with_model()
 * llama_n_ctx_train()
 * llama_n_ctx()
 * llama_load_session_file() 读取保存的session
 * llama_set_rng_seed()
 * llama_should_add_bos_token()
 * llama_tokenize()
 * llama_kv_cache_seq_rm()
 * llama_token_to_piece()
 * llama_sampling_print()
 * llama_sampling_order_print()
 * llama_sampling_init()
 * llama_kv_cache_seq_rm()
 * llama_kv_cache_seq_shift()
 * llama_kv_cache_seq_div()
 * llama_decode() <- llama_eval()
 * llama_batch_get_one()
 * llama_save_session_file()
 * llama_sampling_sample()
 * llama_sampling_accept()
 * llama_sampling_prev_str()
 * llama_sampling_last()
 * llama_sampling_reset()
 * llama_print_timings()
 * llama_free()
 * llama_free_model()
 * llama_sampling_free()
 * llama_backend_free()
 */

extern (C)
{
    enum LlamaModelKvOverrideType
    {
        LLAMA_KV_OVERRIDE_INT,
        LLAMA_KV_OVERRIDE_FLOAT,
        LLAMA_KV_OVERRIDE_BOOL,
    }

    enum LlamaSplitMode
    {
        LLAMA_SPLIT_NONE = 0, // single GPU
        LLAMA_SPLIT_LAYER = 1, // split layers and KV across GPUs
        LLAMA_SPLIT_ROW = 2, // split rows across GPUs
    }

    struct LlamaModelKvOverride
    {
        char[128] key;
        LlamaModelKvOverrideType tag = LlamaModelKvOverrideType.LLAMA_KV_OVERRIDE_INT;
        union
        {
            long int_value;
            double float_value;
            bool bool_value;
        }
    }

    alias llama_progress_callback = bool function(float progress, void* ctx);

    struct LlamaModelParams
    {
        int n_gpu_layers; // number of layers to store in VRAM
        LlamaSplitMode split_mode = LlamaSplitMode.LLAMA_SPLIT_NONE; // how to split the model across multiple GPUs

        // main_gpu interpretation depends on split_mode:
        // LLAMA_SPLIT_NONE: the GPU that is used for the entire model
        // LLAMA_SPLIT_ROW: the GPU that is used for small tensors and intermediate results
        // LLAMA_SPLIT_LAYER: ignored
        int main_gpu;

        // proportion of the model (layers or rows) to offload to each GPU, size: llama_max_devices()
        const float* tensor_split;

        // Called with a progress value between 0.0 and 1.0. Pass NULL to disable.
        // If the provided progress_callback returns true, model loading continues.
        // If it returns false, model loading is immediately aborted.
        llama_progress_callback progress_callback;

        // context pointer passed to the progress callback
        void* progress_callback_user_data;

        // override key-value pairs of the model meta data
        LlamaModelKvOverride* kv_overrides;

        // Keep the booleans together to avoid misalignment during copy-by-value.
        bool vocab_only; // only load the vocabulary, no weights
        bool use_mmap; // use mmap if possible
        bool use_mlock; // force system to keep model in RAM
    }

    enum GgmlType
    {
        GGML_TYPE_F32 = 0,
        GGML_TYPE_F16 = 1,
        GGML_TYPE_Q4_0 = 2,
        GGML_TYPE_Q4_1 = 3,
        // GGML_TYPE_Q4_2 = 4, support has been removed
        // GGML_TYPE_Q4_3 (5) support has been removed
        GGML_TYPE_Q5_0 = 6,
        GGML_TYPE_Q5_1 = 7,
        GGML_TYPE_Q8_0 = 8,
        GGML_TYPE_Q8_1 = 9,
        // k-quantizations
        GGML_TYPE_Q2_K = 10,
        GGML_TYPE_Q3_K = 11,
        GGML_TYPE_Q4_K = 12,
        GGML_TYPE_Q5_K = 13,
        GGML_TYPE_Q6_K = 14,
        GGML_TYPE_Q8_K = 15,
        GGML_TYPE_IQ2_XXS = 16,
        GGML_TYPE_IQ2_XS = 17,
        GGML_TYPE_IQ3_XXS = 18,
        GGML_TYPE_IQ1_S = 19,
        GGML_TYPE_I8,
        GGML_TYPE_I16,
        GGML_TYPE_I32,
        GGML_TYPE_COUNT,
    }

    alias ggml_backend_sched_eval_callback = bool function(void* t, bool ask, void* user_data);

    // (4294967295, 512, 512, 4, 4, -1,
    // 0, 0, -1, 1, 32, 1, 0,
    // null, null,
    // GGML_TYPE_F16, GGML_TYPE_F16,
    // true, false, false, true, true)
    struct LlamaContextParams
    {
        uint seed; // RNG seed, -1 for random
        uint n_ctx; // text context, 0 = from model
        uint n_batch; // prompt processing maximum batch size
        uint n_threads; // number of threads to use for generation
        uint n_threads_batch; // number of threads to use for batch processing
        int rope_scaling_type; // RoPE scaling type, from `enum llama_rope_scaling_type`

        // ref: https://github.com/ggerganov/llama.cpp/pull/2054
        float rope_freq_base; // RoPE base frequency, 0 = from model
        float rope_freq_scale; // RoPE frequency scaling factor, 0 = from model
        float yarn_ext_factor; // YaRN extrapolation mix factor, negative = from model
        float yarn_attn_factor; // YaRN magnitude scaling factor
        float yarn_beta_fast; // YaRN low correction dim
        float yarn_beta_slow; // YaRN high correction dim
        uint yarn_orig_ctx; // YaRN original context size

        ggml_backend_sched_eval_callback cb_eval;
        void* cb_eval_user_data;

        GgmlType type_k = GgmlType.GGML_TYPE_Q2_K; // data type for K cache
        GgmlType type_v = GgmlType.GGML_TYPE_Q2_K; // data type for V cache

        // Keep the booleans together to avoid misalignment during copy-by-value.
        bool mul_mat_q; // if true, use experimental mul_mat_q kernels (DEPRECATED - always true)
        bool logits_all; // the llama_eval() call computes all logits, not just the last one (DEPRECATED - set llama_batch.logits instead)
        bool embedding; // embedding mode only
        bool offload_kqv; // whether to offload the KQV ops (including the KV cache) to GPU
        bool do_pooling; // whether to pool (sum) embedding results by sequence id (ignored if no pooling layer)
    }

    struct LlamaModel;
    struct LlamaContext;

    alias int32_t = int;
    alias uint32_t = uint;

    alias llama_pos = int32_t;
    alias llama_token = int32_t;
    alias llama_seq_id = int32_t;

    // Input data for llama_decode
    // A llama_batch object can contain input about one or many sequences
    // The provided arrays (i.e. token, embd, pos, etc.) must have size of n_tokens
    //
    // - token  : the token ids of the input (used when embd is NULL)
    // - embd   : token embeddings (i.e. float vector of size n_embd) (used when token is NULL)
    // - pos    : the positions of the respective token in the sequence
    // - seq_id : the sequence to which the respective token belongs
    // - logits : if zero, the logits for the respective token will not be output
    //
    struct LlamaBatch
    {
        int32_t n_tokens;

        llama_token* token;
        float* embd;
        llama_pos* pos;
        int32_t* n_seq_id;
        llama_seq_id** seq_id;
        byte* logits;

        // NOTE: helpers for smooth API transition - can be deprecated in the future
        //       for future-proof code, use the above fields instead and ignore everything below
        //
        // pos[i] = all_pos_0 + i*all_pos_1
        //
        llama_pos all_pos_0; // used if pos == NULL
        llama_pos all_pos_1; // used if pos == NULL
        llama_seq_id all_seq_id; // used if seq_id == NULL
    }

    // 下面是一些最简单的会用到的函数，逐步完善丰富

    void llama_backend_init();
    LlamaModelParams llama_model_default_params(); // 加载模型时用到
    LlamaContextParams llama_context_default_params(); // 在create_completion的时候会用到
    LlamaModel* llama_load_model_from_file(const char* path_model, LlamaModelParams params); // 模型支持
    LlamaContext* llama_new_context_with_model(LlamaModel* model, LlamaContextParams params); // 如果需要llama_apply_lora_from_file则需要context
    int llama_model_apply_lora_from_file(
        const LlamaModel* model, const char* path_lora, float scale, const char* path_base_model, int n_threads); // lora支持

    // Allocates a batch of tokens on the heap that can hold a maximum of n_tokens
    // Each token can be assigned up to n_seq_max sequence ids 这里一般设置1就可以
    // The batch has to be freed with llama_batch_free()
    // If embd != 0, llama_batch.embd will be allocated with size of n_tokens * embd * sizeof(float)
    // Otherwise, llama_batch.token will be allocated to store n_tokens llama_token
    // The rest of the llama_batch members are allocated with size n_tokens
    // All members are left uninitialized
    LlamaBatch llama_batch_init(int32_t n_tokens, int32_t embd, int32_t n_seq_max);
    void llama_batch_free(LlamaBatch batch); // Frees a batch of tokens allocated with llama_batch_init()

    // 参考llama_cpp_python调用链条 create_completion -> generate -> eval，最终是调用了_LlamaContext.decode
    int llama_decode(LlamaContext* ctx, LlamaBatch batch);
    void llama_free(LlamaContext* ctx); // 释放context
    void llama_free_model(LlamaModel* model); // 释放模型
    void llama_backend_free();

    /// @details Convert the provided text into tokens.
    /// @param tokens The tokens pointer must be large enough to hold the resulting tokens.
    /// @return Returns the number of tokens on success, no more than n_max_tokens
    /// @return Returns a negative number on failure - the number of tokens that would have been returned
    /// @param special Allow tokenizing special and/or control tokens which otherwise are not exposed and treated as plaintext.
    ///                Does not insert a leading space.
    int32_t llama_tokenize(
        const LlamaModel* model, // 模型指针
        const char* text, // prompt pointer
        int32_t text_len, // prompt buffer size
        llama_token* tokens, // token pointer
        int32_t n_max_tokens, // token pointer max buffer size 当不够空间的时候，需要重新分配
        bool add_bos, // 是否添加bos标示
        bool special);

    // Token Id -> Piece.
    // Uses the vocabulary in the provided context.
    // Does not write null terminator to the buffer.
    // User code is responsible to remove the leading whitespace of the first non-BOS token when decoding multiple tokens.
    int32_t llama_token_to_piece(
        const LlamaModel* model,
        llama_token token,
        char* buf,
        int32_t length);

    alias llama_detokenize = llama_token_to_piece;

    uint32_t llama_n_ctx(const LlamaContext* ctx);
    uint32_t llama_n_batch(const LlamaContext* ctx);
    float* llama_get_logits_ith(LlamaContext* ctx, int32_t i);
    int32_t llama_n_vocab(const LlamaModel* model);

    llama_token llama_token_bos(const LlamaModel* model); // beginning-of-sentence
    llama_token llama_token_eos(const LlamaModel* model); // end-of-sentence
    llama_token llama_token_nl(const LlamaModel* model); // next-line

    struct LlamaTokenData
    {
        llama_token id; // token id
        float logit; // log-odds of the token
        float p; // probability of the token
    }

    struct LlamaTokenDataArray
    {
        LlamaTokenData* data;
        size_t size;
        bool sorted;
    }

    llama_token llama_sample_token_greedy(LlamaContext* ctx, LlamaTokenDataArray* candidates);

    // 需要实现 llama_batch_add()
    // void llama_batch_add(
    //             struct llama_batch & batch,
    //                    llama_token   id,
    //                      llama_pos   pos,
    // const std::vector<llama_seq_id> & seq_ids,
    //                            bool   logits) {
    //     batch.token   [batch.n_tokens] = id;
    //     batch.pos     [batch.n_tokens] = pos;
    //     batch.n_seq_id[batch.n_tokens] = seq_ids.size();
    //     for (size_t i = 0; i < seq_ids.size(); ++i) {
    //         batch.seq_id[batch.n_tokens][i] = seq_ids[i];
    //     }
    //     batch.logits  [batch.n_tokens] = logits;
    //
    //     batch.n_tokens++;
    // }
    void llama_batch_add(ref LlamaBatch batch, llama_token id, llama_pos pos, int[] seq_ids, bool logits)
    {
        batch.token[batch.n_tokens] = id;
        batch.pos[batch.n_tokens] = pos;
        batch.n_seq_id[batch.n_tokens] = cast(int) seq_ids.length;
        for (int i = 0; i < seq_ids.length; i++)
        {
            batch.seq_id[batch.n_tokens][i] = seq_ids[i];
        }
        batch.logits[batch.n_tokens] = logits;
        batch.n_tokens++;
    }

    void llama_batch_clear(ref LlamaBatch batch)
    {
        batch.n_tokens = 0;
    }
}

unittest
{
    // 这里尝试跑通一个例子，作为编程实现的基础，参考llama.cpp examples/simple/simple.cpp
    // llama.cpp为了适配底层的模型，以及C语言，在底层实现的

    import std.stdio;

    llama_backend_init();

    LlamaModelParams params = llama_model_default_params();
    LlamaContextParams ctxParams = llama_context_default_params();
    // 一些看起来要设置的参数，跟硬件环境相关
    ctxParams.n_threads = 2;
    ctxParams.n_threads_batch = 2;

    LlamaModel* llamaModel = llama_load_model_from_file("model/llama-2-7b-chat.Q2_K.gguf", params);
    LlamaContext* llamaContext = llama_new_context_with_model(llamaModel, ctxParams);

    if (llamaContext == null)
    {
        writefln("failed to init context");
        return;
    }

    // 按需加载lora文件
    // if (llama_model_apply_lora_from_file(llamaModel,
    //         lora_path, // 指向lora文件的路径
    //         1.0, // lora_scale,
    //         null, // lora_base设置为null，从而使用当前model
    //         ctxParams.n_threads) != 0)
    // {
    //     // 异常退出处理
    // }

    char[] prompt = "hello".dup();

    // 将prompt转成token
    llama_token[] tokens_list = new llama_token[ctxParams.n_batch];
    // llama_tokenize
    writeln("prompt<" ~ prompt ~ ">length:", prompt.length);
    writeln("tokenize:", llama_tokenize(llamaModel, &prompt[0], cast(int) prompt.length, &tokens_list[0], ctxParams
            .n_batch, true, false));

    // 检查kv cache size是否满足，与n_ctx比较
    // 这里要特别注意检查：n_kv_req, n_ctx, n_len之间的关系
    // n_kv_req = prompt + input
    // n_len = input
    // n_ctx 是允许输入的最大长度
    if (ctxParams.n_ctx < prompt.length)
    {
        writefln("n_ctx not enough");
        return;
    }

    // 默认情况下，n_ctx == n_batch
    // 初始化batch
    LlamaBatch batch = llama_batch_init(ctxParams.n_batch, 0, 1); // embd = 0，目前不支持设置

    writeln("batch for debug:", batch);

    // llama_batch_add 做了什么有什么作用？把prompt添加到batch中
    // evaluate the initial prompt
    // for (size_t i = 0; i < tokens_list.size(); i++) {
    //    llama_batch_add(batch, tokens_list[i], i, { 0 }, false);
    // }
    writeln("batch n_token for debug:", batch.n_tokens);
    for (int i = 0; i < tokens_list.length; i++)
    {
        if (tokens_list[i] != 0)
        {
            llama_batch_add(batch, tokens_list[i], i, [0], false);
        }
        else
            break;
    }
    // llama_decode will output logits only for the last token of the prompt
    batch.logits[batch.n_tokens - 1] = true;

    writeln("llama_decode for debug(should be 0):", llama_decode(llamaContext, batch));
    writeln(">>", prompt);

    // while循环持续处理 sample & eval
    // 1. sample the next token, llama_sample_token_greedy()
    // 2. if (new_token_id == llama_token_eos(model) || n_cur == n_len) break;
    // 3. llama_batch_clear() // 需要自己实现
    // 4. llama_batch_add() // 需要自己实现
    // 5. llama_decode()
    int n_cur = cast(int) prompt.length; // batch.n_tokens 这里是指向添加了token_list后batch.n_tokens的当前位置

    int max_len = ctxParams.n_batch / 2;
    while (n_cur <= max_len)
    {
        // sample the next token
        int n_vocab = llama_n_vocab(llamaModel);
        float* logits = llama_get_logits_ith(llamaContext, batch.n_tokens - 1);

        LlamaTokenData[] candidates = new LlamaTokenData[n_vocab];
        for (llama_token token_id = 0; token_id < n_vocab; token_id++)
        {
            candidates[token_id].id = token_id;
            candidates[token_id].logit = logits[token_id];
            candidates[token_id].p = 0.0f;
        }

        LlamaTokenDataArray candidates_arr = {
            &candidates[0], candidates.length, false
        };
        const llama_token new_token_id = llama_sample_token_greedy(llamaContext, &candidates_arr);
        if (new_token_id == llama_token_eos(llamaModel))
        {
            writeln("eos:", new_token_id);
            break;
        }

        {
            char[32] result;
            int n_tokens = llama_detokenize(llamaModel, new_token_id, &result[0], result.length);
            assert(n_tokens > 0 && n_tokens < 32);

            {
                char[] result2 = new char[n_tokens];
                for (int i = 0; i < n_tokens; i++)
                {
                    result2[i] = result[i];
                }
                write(result2);
                stdout.flush();
            }
        }

        llama_batch_clear(batch);
        llama_batch_add(batch, new_token_id, n_cur, [0], true);

        n_cur++;
        if (n_cur >= max_len)
        {
            writeln("max_len:", n_cur, max_len);
            break;
        }

        // eval
        if (llama_decode(llamaContext, batch))
        {
            writeln("exception when llama_decode()");
            break;
        }
    }

    // 输出最终结果

    llama_batch_free(batch);

    llama_free(llamaContext);
    llama_free_model(llamaModel);
    llama_backend_free();
}
